package com.kexio.auth.async;

import com.kexio.auth.context.TenantContextHolder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.slf4j.MDC;
import org.springframework.core.task.TaskDecorator;
import org.springframework.util.StringUtils;

import java.util.Map;

/**
 * 租户感知的任务装饰器
 * 
 * 用于在异步任务执行时传递和管理租户上下文：
 * 1. 在任务提交时：捕获当前线程的租户上下文
 * 2. 在任务执行时：在异步线程中恢复租户上下文
 * 3. 在任务完成时：清理异步线程的租户上下文
 * 4. 异常处理：确保即使任务执行失败也能正确清理上下文
 * 
 * 支持场景：
 * - @Async 异步方法调用
 * - ThreadPoolTaskExecutor 线程池任务
 * - CompletableFuture 异步编程
 * - 定时任务（@Scheduled）
 * 
 * 使用方式：
 * <pre>
 * &#64;Configuration
 * public class AsyncConfig {
 *     &#64;Bean
 *     public TaskExecutor taskExecutor() {
 *         ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor();
 *         executor.setTaskDecorator(new TenantAwareTaskDecorator());
 *         return executor;
 *     }
 * }
 * </pre>
 * 
 * @author Kexio Team
 * @since 1.0.0
 */
public class TenantAwareTaskDecorator implements TaskDecorator {
    
    private static final Logger logger = LoggerFactory.getLogger(TenantAwareTaskDecorator.class);
    
    /**
     * MDC中租户ID的键名
     */
    private static final String MDC_TENANT_KEY = "tenantId";
    
    /**
     * 是否启用调试模式
     */
    private final boolean debugMode;
    
    /**
     * 构造函数
     */
    public TenantAwareTaskDecorator() {
        this(false);
    }
    
    /**
     * 构造函数
     * 
     * @param debugMode 是否启用调试模式
     */
    public TenantAwareTaskDecorator(boolean debugMode) {
        this.debugMode = debugMode;
    }
    
    @Override
    public Runnable decorate(Runnable runnable) {
        // 1. 捕获当前线程的租户上下文
        String currentTenantId = captureTenantContext();
        Map<String, String> currentMDC = captureMDCContext();
        
        if (debugMode) {
            logger.debug("装饰异步任务，捕获租户上下文: {} (任务: {})", 
                currentTenantId, runnable.getClass().getSimpleName());
        }
        
        // 2. 返回包装后的任务
        return new TenantAwareRunnable(runnable, currentTenantId, currentMDC, debugMode);
    }
    
    /**
     * 捕获当前线程的租户上下文
     */
    private String captureTenantContext() {
        try {
            return TenantContextHolder.getTenantOrNull();
        } catch (Exception e) {
            logger.debug("捕获租户上下文失败: {}", e.getMessage());
            return null;
        }
    }
    
    /**
     * 捕获当前线程的MDC上下文
     */
    private Map<String, String> captureMDCContext() {
        try {
            return MDC.getCopyOfContextMap();
        } catch (Exception e) {
            logger.debug("捕获MDC上下文失败: {}", e.getMessage());
            return null;
        }
    }
    
    /**
     * 租户感知的Runnable包装器
     */
    private static class TenantAwareRunnable implements Runnable {
        
        private static final Logger logger = LoggerFactory.getLogger(TenantAwareRunnable.class);
        
        private final Runnable delegate;
        private final String tenantId;
        private final Map<String, String> mdcContext;
        private final boolean debugMode;
        
        public TenantAwareRunnable(Runnable delegate, String tenantId, 
                                 Map<String, String> mdcContext, boolean debugMode) {
            this.delegate = delegate;
            this.tenantId = tenantId;
            this.mdcContext = mdcContext;
            this.debugMode = debugMode;
        }
        
        @Override
        public void run() {
            String originalTenantId = null;
            Map<String, String> originalMDC = null;
            
            try {
                // 1. 备份异步线程的原始上下文（如果有的话）
                originalTenantId = TenantContextHolder.getTenantOrNull();
                originalMDC = MDC.getCopyOfContextMap();
                
                // 2. 设置从主线程传递过来的租户上下文
                if (StringUtils.hasText(tenantId)) {
                    TenantContextHolder.setTenant(tenantId);
                    if (debugMode) {
                        logger.debug("异步任务中设置租户上下文: {} (任务: {})", 
                            tenantId, delegate.getClass().getSimpleName());
                    }
                }
                
                // 3. 设置MDC上下文
                if (mdcContext != null && !mdcContext.isEmpty()) {
                    MDC.setContextMap(mdcContext);
                } else if (StringUtils.hasText(tenantId)) {
                    // 如果没有完整的MDC上下文，至少设置租户ID
                    MDC.put(MDC_TENANT_KEY, tenantId);
                }
                
                // 4. 执行原始任务
                delegate.run();
                
            } catch (Exception e) {
                logger.error("异步任务执行发生异常: {} (租户: {}, 任务: {})", 
                    e.getMessage(), tenantId, delegate.getClass().getSimpleName(), e);
                throw e;
            } finally {
                // 5. 清理或恢复上下文（确保总是执行）
                cleanupOrRestoreContext(originalTenantId, originalMDC);
            }
        }
        
        /**
         * 清理或恢复上下文
         */
        private void cleanupOrRestoreContext(String originalTenantId, Map<String, String> originalMDC) {
            try {
                // 清理当前设置的租户上下文
                TenantContextHolder.clear();
                
                // 恢复原始的租户上下文（如果有的话）
                if (StringUtils.hasText(originalTenantId)) {
                    TenantContextHolder.setTenant(originalTenantId);
                }
                
                // 清理MDC
                MDC.clear();
                
                // 恢复原始的MDC上下文（如果有的话）
                if (originalMDC != null && !originalMDC.isEmpty()) {
                    MDC.setContextMap(originalMDC);
                }
                
                if (debugMode) {
                    logger.debug("异步任务上下文清理完成: {} -> {} (任务: {})", 
                        tenantId, originalTenantId, delegate.getClass().getSimpleName());
                }
                
            } catch (Exception e) {
                logger.error("异步任务上下文清理失败: {} (任务: {})", 
                    e.getMessage(), delegate.getClass().getSimpleName(), e);
                
                // 强制清理，防止上下文泄露
                try {
                    TenantContextHolder.clear();
                    MDC.clear();
                } catch (Exception forceCleanupException) {
                    logger.error("强制清理异步任务上下文失败: {}", forceCleanupException.getMessage());
                }
            }
        }
        
        @Override
        public String toString() {
            return String.format("TenantAwareRunnable{tenant='%s', delegate=%s}", 
                tenantId, delegate.getClass().getSimpleName());
        }
    }
    
    // ==================== 工厂方法 ====================
    
    /**
     * 创建默认的租户感知任务装饰器
     */
    public static TenantAwareTaskDecorator create() {
        return new TenantAwareTaskDecorator();
    }
    
    /**
     * 创建带调试模式的租户感知任务装饰器
     */
    public static TenantAwareTaskDecorator createWithDebug() {
        return new TenantAwareTaskDecorator(true);
    }
    
    // ==================== 监控方法 ====================
    
    /**
     * 获取装饰器信息
     */
    public DecoratorInfo getInfo() {
        return new DecoratorInfo();
    }
    
    /**
     * 装饰器信息
     */
    public class DecoratorInfo {
        public String getName() {
            return "TenantAwareTaskDecorator";
        }
        
        public String getDescription() {
            return "租户感知的任务装饰器 - 支持异步任务中的租户上下文传递";
        }
        
        public boolean isDebugMode() {
            return debugMode;
        }
        
        public String getVersion() {
            return "1.0.0";
        }
        
        @Override
        public String toString() {
            return String.format("DecoratorInfo{name='%s', debugMode=%s, version='%s'}", 
                getName(), debugMode, getVersion());
        }
    }
}
